from typing import List, Dict
from sklearn.neighbors import KDTree
import numpy as np
import asyncio
import time
from model_server import model_predict
import json

def task_allocation_kdtree(model,task_points,drone_points):
    """
    使用 KDTree 将任务点分配给最近的无人机起点。

    Args:
        data (RequestData): 包含任务点和无人机起点的数据对象。

    Returns:
        Dict[str, List[List[float]]]: 返回的字典，键是 "Drone 1", "Drone 2" 等，
                                       值是对应无人机及其负责的任务点。
    """


    drone_points_np = np.array(drone_points)
    # 构建 KDTree
    kdtree = KDTree(drone_points_np)
    task_points_np = np.array(task_points)
    
    if len(drone_points_np) == 0:
        raise ValueError("无人机起点列表不能为空")
    if len(task_points_np) == 0:
        raise ValueError("任务点列表不能为空")

    if drone_points_np.ndim != 2 or drone_points_np.shape[1] != 2:
        raise ValueError("无人机起点必须是二维坐标")

    if task_points_np.ndim != 2 or task_points_np.shape[1] != 2:
        raise ValueError("任务点必须是二维坐标")
    
    # 查询每个任务点最近的无人机索引
    _, nearest_indices = kdtree.query(task_points_np)

    drone_res = {f"Drone {i+1}": [] for i in range(len(drone_points))}
        
    for index in range(len(drone_points)):
        indexs_xy = np.where(nearest_indices == index)[0]
        agent_index = f"Drone {index+1}"
        drone_res.setdefault(agent_index, []).append(drone_points[index])
        for idx in indexs_xy:
            drone_res.setdefault(agent_index, []).append(task_points[idx])
        drone_res[agent_index] = model_predict(model, drone_res[agent_index], drone_points[index])
    
    # 初始排序
    # for index in zip(drone_res):

    #     drone_res[key] = asyncio.run(model_predict(model,drone_res[key], drone_points[index]))


    return drone_res

def calculate_average_position(path):
    if not path:
        return np.array([float('inf')] * len(path[0])) if path else np.array([float('inf')])
    return np.mean(path, axis=0)

def reallocate_tasks(agents_positions,unvisited_points,damaged_agents,model=None):
    print("agents_positions:", agents_positions)
    
    damaged_path = []
    # 获取损毁智能体的路径并清空
    for damaged_agent in damaged_agents:
        print("damaged_agent:",damaged_agent)
        damaged_agent_key = f"Drone {damaged_agent + 1}"
        damaged_path += agents_positions.pop(damaged_agent_key, [])

    if not damaged_path:
        print("No damaged path found for agent.")
        return
    # 找到未被遍历的点
    print("unvisited_points:",unvisited_points)
    start_time = time.time()
    # 找到可用的智能体
    available_agents = [agent for agent in agents_positions.keys() if agent not in damaged_agents]
    print("available_agents:",available_agents)
    if not available_agents:
        print("No available agents to reallocate tasks.")
        return

    # # 遍历未被遍历的点，分配给最近的可用智能体
    for point in unvisited_points:
        closest_agent = min(
            available_agents,
            key=lambda agent: np.linalg.norm(
                calculate_average_position(agents_positions[agent]) - np.array(point)
            ) if agents_positions[agent] else float('inf')
        )
        agents_positions[closest_agent].append(point)
    return agents_positions,(time.time() - start_time)


# if __name__ == "__main__":
#     # 示例任务点和无人机起点
#     task_points = [
#         [0.0, 0.0],
#         [0.2, 0.5],
#         [0.7, 0.3],
#         [1.0, 1.0],
#         [0.5, 0.5]
#     ]
    
#     drone_points = [
#         [0.0, 0.0],
#         [1.0, 1.0]
#     ]
#     from utils import load_model 
#     model = load_model("params/tsp_100",device=0)[0]
#     # 调用 task_allocation_kdtree 进行测试
#     result = task_allocation_kdtree(model,task_points, drone_points)
    
#     # 打印结果
#     print("Task Allocation Result:")
#     for agent, points in result.items():
#         print(f"{agent}: {points}")
if __name__ == "__main__":
    # 模拟无人机路径信息
    agents_positions = {
        "Drone 1": [[0.0, 0.0], [0.2, 0.5], [0.7, 0.3]],
        "Drone 2": [[1.0, 1.0], [0.5, 0.5]]
    }

    # 模拟网格点及其覆盖状态（第三个元素表示是否已访问）
    unvisited_points = [[0.2, 0.5], [0.7, 0.3],[0.5, 0.5]]

    # 模拟损毁的无人机索引列表
    damaged_agents = [0]  # 假设 Drone 1 (index 0) 被损毁

    # 模拟模型加载
    from utils import load_model
    model = load_model("params/tsp_100", device=0)[0]

    # 调用 reallocate_tasks 进行测试
    execution_time = reallocate_tasks(agents_positions, unvisited_points, damaged_agents, model)

    # 打印更新后的任务分配结果
    print("Updated Task Allocation after Reallocation:")
    for agent, points in agents_positions.items():
        print(f"{agent}: {points}")

    # 打印执行时间
    print(f"Reallocate tasks executed in {execution_time:.4f} seconds")